from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
from sklearn.metrics import adjusted_rand_score, silhouette_score


# 定义KMeans模型
class KMeansModel:
    def __init__(self, n_clusters=3):
        self.n_clusters = n_clusters
        self.kmeans = KMeans(n_clusters=n_clusters, n_init=100)

    def fit(self, X):
        self.kmeans.fit(X)
        return self

    def predict(self, X):
        return self.kmeans.predict(X)


if __name__ == '__main__':
    # 生成数据集
    data = make_blobs(n_samples=100, centers=4, random_state=42)
    X = data[0]
    true_labels = data[1]

    # 训练模型
    model = KMeansModel()
    kmeans_model = model.fit(X)
    pred_labels = model.predict(X)
    print("pred_labels: ", pred_labels)
    # 计算ARI和轮廓系数
    ari = adjusted_rand_score(true_labels, pred_labels)
    print("ARI: ", ari)
    silhouette = silhouette_score(X, true_labels)
    print("Silhouette Coefficient: ", silhouette)
